{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ISubpr_SSsiM"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3jTMb1dySr3V"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6DWfyNThSziV"
      },
      "source": [
        "# Better performance with tf.function\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/function\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J122XQYG7W6w"
      },
      "source": [
        "In TensorFlow 2, eager execution is turned on by default. The user interface is intuitive and flexible (running one-off operations is much easier\n",
        "and faster), but this can come at the expense of performance and deployability.\n",
        "\n",
        "You can use `tf.function` to make graphs out of your programs. It is a transformation tool that creates Python-independent dataflow graphs out of your Python code. This will help you create performant and portable models, and it is required to use `SavedModel`.\n",
        "\n",
        "This guide will help you conceptualize how `tf.function` works under the hood so you can use it effectively.\n",
        "\n",
        "The main takeaways and recommendations are:\n",
        "\n",
        "- Debug in eager mode, then decorate with `@tf.function`.\n",
        "- Don't rely on Python side effects like object mutation or list appends.\n",
        "- `tf.function` works best with TensorFlow ops; NumPy and Python calls are converted to constants.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SjvqpgepHJPd"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "otIdN1TS8N7S"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0xDjO4SHLUD"
      },
      "source": [
        "Define a helper function to demonstrate the kinds of errors you might encounter:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D25apou9IOXa"
      },
      "outputs": [],
      "source": [
        "import traceback\n",
        "import contextlib\n",
        "\n",
        "# Some helper code to demonstrate the kinds of errors you might encounter.\n",
        "@contextlib.contextmanager\n",
        "def assert_raises(error_class):\n",
        "  try:\n",
        "    yield\n",
        "  except error_class as e:\n",
        "    print('Caught expected exception \\n  {}:'.format(error_class))\n",
        "    traceback.print_exc(limit=2)\n",
        "  except Exception as e:\n",
        "    raise e\n",
        "  else:\n",
        "    raise Exception('Expected {} to be raised but no error was raised!'.format(\n",
        "        error_class))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WPSfepzTHThq"
      },
      "source": [
        "## Basics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CNwYTIJ8r56W"
      },
      "source": [
        "### Usage\n",
        "\n",
        "A `Function` you define (for example by applying the `@tf.function` decorator) is just like a core TensorFlow operation: You can execute it eagerly; you can compute gradients; and so on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SbtT1-Wm70F2"
      },
      "outputs": [],
      "source": [
        "@tf.function  # The decorator converts `add` into a `Function`.\n",
        "def add(a, b):\n",
        "  return a + b\n",
        "\n",
        "add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uP-zUelB8DbX"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(1.0)\n",
        "with tf.GradientTape() as tape:\n",
        "  result = add(v, 1.0)\n",
        "tape.gradient(result, v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocWZvqrmHnmX"
      },
      "source": [
        "You can use `Function`s inside other `Function`s."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l5qRjdbBVdU6"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def dense_layer(x, w, b):\n",
        "  return add(tf.matmul(x, w), b)\n",
        "\n",
        "dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "piBhz7gYsHqU"
      },
      "source": [
        "`Function`s can be faster than eager code, especially for graphs with many small ops. But for graphs with a few expensive ops (like convolutions), you may not see much speedup.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zuXt4wRysI03"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uZ4Do2AV80cO"
      },
      "source": [
        "### Tracing\n",
        "\n",
        "This section exposes how `Function` works under the hood, including implementation details *which may change in the future*. However, once you understand why and when tracing happens, it's much easier to use `tf.function` effectively!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nhpUtRqsXoyM"
      },
      "source": [
        "#### What is \"tracing\"?\n",
        "\n",
        "A `Function` runs your program in a [TensorFlow Graph](https://www.tensorflow.org/guide/intro_to_graphs#what_are_graphs). However, a `tf.Graph` cannot represent all the things that you'd write in an eager TensorFlow program. For instance, Python supports polymorphism, but `tf.Graph` requires its inputs to have a specified data type and dimension. Or you may perform side tasks like reading command-line arguments, raising an error, or working with a more complex Python object; none of these things can run in a `tf.Graph`.\n",
        "\n",
        "`Function` bridges this gap by separating your code in two stages:\n",
        "\n",
        "  1)  In the first stage, referred to as \"**tracing**\", `Function` creates a new `tf.Graph`. Python code runs normally, but all TensorFlow operations (like adding two Tensors) are *deferred*: they are captured by the `tf.Graph` and not run.\n",
        "\n",
        "  2) In the second stage, a `tf.Graph` which contains everything that was deferred in the first stage is run. This stage is much faster than the tracing stage.\n",
        "\n",
        "Depending on its inputs, `Function` will not always run the first stage when it is called.  See [\"Rules of tracing\"](#rules_of_tracing) below to get a better sense of how it makes that determination. Skipping the first stage and only executing the second stage is what gives you TensorFlow's high performance.\n",
        "\n",
        "When `Function` does decide to trace, the tracing stage is immediately followed by the second stage, so calling the `Function` both creates and runs the `tf.Graph`. Later you will see how you can run only the tracing stage with [`get_concrete_function`](#obtaining_concrete_functions)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K7scSzLx662f"
      },
      "source": [
        "When we pass arguments of different types into a `Function`, both stages are run:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kojmJrgq8U9v"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def double(a):\n",
        "  print(\"Tracing with\", a)\n",
        "  return a + a\n",
        "\n",
        "print(double(tf.constant(1)))\n",
        "print()\n",
        "print(double(tf.constant(1.1)))\n",
        "print()\n",
        "print(double(tf.constant(\"a\")))\n",
        "print()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPfouGUQrcNb"
      },
      "source": [
        "Note that if you repeatedly call a `Function` with the same argument type, TensorFlow will skip the tracing stage and reuse a previously traced graph, as the generated graph would be identical."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hFccbWFRrsBp"
      },
      "outputs": [],
      "source": [
        "# This doesn't print 'Tracing with ...'\n",
        "print(double(tf.constant(\"b\")))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fgIO_XEzcB9o"
      },
      "source": [
        "You can use `pretty_printed_concrete_signatures()` to see all of the available traces:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IiQc4IKAb-NX"
      },
      "outputs": [],
      "source": [
        "print(double.pretty_printed_concrete_signatures())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rKQ92VEWI7n8"
      },
      "source": [
        "So far, you've seen that `tf.function` creates a cached, dynamic dispatch layer over TensorFlow's graph tracing logic. To be more specific about the terminology:\n",
        "\n",
        "- A `tf.Graph` is the raw, language-agnostic, portable representation of a TensorFlow computation.\n",
        "- A `ConcreteFunction` wraps a `tf.Graph`.\n",
        "- A `Function` manages a cache of `ConcreteFunction`s and picks the right one for your inputs.\n",
        "- `tf.function` wraps a Python function, returning a `Function` object.\n",
        "- **Tracing** creates a `tf.Graph` and wraps it in a `ConcreteFunction`, also known as a **trace.**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "129-iRsPS-gY"
      },
      "source": [
        "#### Rules of tracing\n",
        "\n",
        "A `Function` determines whether to reuse a traced `ConcreteFunction` by computing a **cache key** from an input's args and kwargs. A **cache key** is a key that identifies a `ConcreteFunction` based on the input args and kwargs of the `Function` call, according to the following rules (which may change):\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h62XoXho6EWN"
      },
      "source": [
        "\n",
        "\n",
        "- The key generated for a `tf.Tensor` argument is its shape and dtype.\n",
        "- Starting in TensorFlow 2.3, the key generated for a `tf.Variable` argument is its `id()`.\n",
        "- The key generated for a Python primitive is its value. The key generated for nested `dict`s, `list`s, `tuple`s, `namedtuple`s, and [`attr`](https://www.attrs.org/en/stable/)s is the flattened tuple. (As a result of this flattening, calling a concrete function with a different nesting structure than the one used during tracing will result in a TypeError).\n",
        "- For all other Python types, the keys are based on the object `id()` so that methods are traced independently for each instance of a class.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GNNN4lgRzpIs"
      },
      "source": [
        "\n",
        "Note: Cache keys are based on the `Function` input parameters so changes to global and [free variables](https://docs.python.org/3/reference/executionmodel.html#binding-of-names) alone will not create a new trace. See [this section](#depending_on_python_global_and_free_variables) for recommended practices when dealing with Python global and free variables."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PEDwbumO32Wh"
      },
      "source": [
        "#### Controlling retracing\n",
        "\n",
        "Retracing, which is when your `Function` creates more than one trace, helps ensures that TensorFlow generates correct graphs for each set of inputs. However, tracing is an expensive operation! If your `Function` retraces a new graph for every call, you'll find that your code executes more slowly than if you didn't use `tf.function`.\n",
        "\n",
        "To control the tracing behavior, you can use the following techniques:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EUtycWJa34TT"
      },
      "source": [
        "- Specify `input_signature` in `tf.function` to limit tracing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_BDMIRmu1RGB"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))\n",
        "def next_collatz(x):\n",
        "  print(\"Tracing with\", x)\n",
        "  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)\n",
        "\n",
        "print(next_collatz(tf.constant([1, 2])))\n",
        "# We specified a 1-D tensor in the input signature, so this should fail.\n",
        "with assert_raises(ValueError):\n",
        "  next_collatz(tf.constant([[1, 2], [3, 4]]))\n",
        "\n",
        "# We specified an int32 dtype in the input signature, so this should fail.\n",
        "with assert_raises(ValueError):\n",
        "  next_collatz(tf.constant([1.0, 2.0]))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocxX-HVk7P2o"
      },
      "source": [
        "- Specify a \\[None\\] dimension in `tf.TensorSpec` to allow for flexibility in trace reuse.\n",
        "\n",
        "  Since TensorFlow matches tensors based on their shape, using a `None` dimension as a wildcard will allow `Function`s to reuse traces for variably-sized input. Variably-sized input can occur if you have sequences of different length, or images of different sizes for each batch (See [Transformer](../tutorials/text/transformer.ipynb) and [Deep Dream](../tutorials/generative/deepdream.ipynb) tutorials for example)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4Viun7dh7PmF"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))\n",
        "def g(x):\n",
        "  print('Tracing with', x)\n",
        "  return x\n",
        "\n",
        "# No retrace!\n",
        "print(g(tf.constant([1, 2, 3])))\n",
        "print(g(tf.constant([1, 2, 3, 4, 5])))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AY5oiQN0XIyA"
      },
      "source": [
        "- Cast Python arguments to Tensors to reduce retracing.\n",
        "\n",
        "  Often, Python arguments are used to control hyperparameters and graph constructions - for example, `num_layers=10` or `training=True` or `nonlinearity='relu'`. So if the Python argument changes, it makes sense that you'd have to retrace the graph.\n",
        "\n",
        "  However, it's possible that a Python argument is not being used to control graph construction. In these cases, a change in the Python value can trigger needless retracing. Take, for example, this training loop, which AutoGraph will dynamically unroll. Despite the multiple traces, the generated graph is actually identical, so retracing is unnecessary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uydzR5JYUU8H"
      },
      "outputs": [],
      "source": [
        "def train_one_step():\n",
        "  pass\n",
        "\n",
        "@tf.function\n",
        "def train(num_steps):\n",
        "  print(\"Tracing with num_steps = \", num_steps)\n",
        "  tf.print(\"Executing with num_steps = \", num_steps)\n",
        "  for _ in tf.range(num_steps):\n",
        "    train_one_step()\n",
        "\n",
        "print(\"Retracing occurs for different Python arguments.\")\n",
        "train(num_steps=10)\n",
        "train(num_steps=20)\n",
        "\n",
        "print()\n",
        "print(\"Traces are reused for Tensor arguments.\")\n",
        "train(num_steps=tf.constant(10))\n",
        "train(num_steps=tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4pJqkDR_Q2wz"
      },
      "source": [
        "If you need to force retracing, create a new `Function`. Separate `Function` objects are guaranteed not to share traces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uHp4ousu4DdN"
      },
      "outputs": [],
      "source": [
        "def f():\n",
        "  print('Tracing!')\n",
        "  tf.print('Executing')\n",
        "\n",
        "tf.function(f)()\n",
        "tf.function(f)()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96IxS2WR37fF"
      },
      "source": [
        "### Obtaining concrete functions\n",
        "\n",
        "Every time a function is traced, a new concrete function is created. You can directly obtain a concrete function, by using `get_concrete_function`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mHg2CGtPQ3Hz"
      },
      "outputs": [],
      "source": [
        "print(\"Obtaining concrete trace\")\n",
        "double_strings = double.get_concrete_function(tf.constant(\"a\"))\n",
        "print(\"Executing traced function\")\n",
        "print(double_strings(tf.constant(\"a\")))\n",
        "print(double_strings(a=tf.constant(\"b\")))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6IVZ-NVf9vsx"
      },
      "outputs": [],
      "source": [
        "# You can also call get_concrete_function on an InputSpec\n",
        "double_strings_from_inputspec = double.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.string))\n",
        "print(double_strings_from_inputspec(tf.constant(\"c\")))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iR4fVmG34xvF"
      },
      "source": [
        "Printing a `ConcreteFunction` displays a summary of its input arguments (with types) and its output type."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o3-JbkIk41r8"
      },
      "outputs": [],
      "source": [
        "print(double_strings)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QtqfvljZeuOV"
      },
      "source": [
        "You can also directly retrieve a concrete function's signature."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nzbrqFABe0zG"
      },
      "outputs": [],
      "source": [
        "print(double_strings.structured_input_signature)\n",
        "print(double_strings.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lar5A_5m5IG1"
      },
      "source": [
        "Using a concrete trace with incompatible types will throw an error"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G5eeTK-T5KYj"
      },
      "outputs": [],
      "source": [
        "with assert_raises(tf.errors.InvalidArgumentError):\n",
        "  double_strings(tf.constant(1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "st2L9VNQVtSG"
      },
      "source": [
        "You may notice that Python arguments are given special treatment in a concrete function's input signature. Prior to TensorFlow 2.3, Python arguments were simply removed from the concrete function's signature. Starting with TensorFlow 2.3, Python arguments remain in the signature, but are constrained to take the value set during tracing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U_QyPSGoaC35"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def pow(a, b):\n",
        "  return a ** b\n",
        "\n",
        "square = pow.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)\n",
        "print(square)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E76vIDhQbXIb"
      },
      "outputs": [],
      "source": [
        "assert square(tf.constant(10.0)) == 100\n",
        "\n",
        "with assert_raises(TypeError):\n",
        "  square(tf.constant(10.0), b=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "41gJh_JGIfuA"
      },
      "source": [
        "### Obtaining graphs\n",
        "\n",
        "Each concrete function is a callable wrapper around a `tf.Graph`. Although retrieving the actual `tf.Graph` object is not something you'll normally need to do, you can obtain it easily from any concrete function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UENeGHfaX8g"
      },
      "outputs": [],
      "source": [
        "graph = double_strings.graph\n",
        "for node in graph.as_graph_def().node:\n",
        "  print(f'{node.input} -> {node.name}')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aIKkgr6qdtp4"
      },
      "source": [
        "### Debugging\n",
        "\n",
        "In general, debugging code is easier in eager mode than inside `tf.function`. You should ensure that your code executes error-free in eager mode before decorating with `tf.function`. To assist in the debugging process, you can call `tf.config.run_functions_eagerly(True)` to globally disable and reenable `tf.function`.\n",
        "\n",
        "When tracking down issues that only appear within `tf.function`, here are some tips:\n",
        "- Plain old Python `print` calls only execute during tracing, helping you track down when your function gets (re)traced.\n",
        "- `tf.print` calls will execute every time, and can help you track down intermediate values during execution.\n",
        "- `tf.debugging.enable_check_numerics` is an easy way to track down where NaNs and Inf are created.\n",
        "- `pdb` can help you understand what's going on during tracing. (Caveat: PDB will drop you into AutoGraph-transformed source code.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f05Vr_YBUCz"
      },
      "source": [
        "## AutoGraph Transformations\n",
        "\n",
        "AutoGraph is a library that is on by default in `tf.function`, and transforms a subset of Python eager code into graph-compatible TensorFlow ops. This includes control flow like `if`, `for`, `while`.\n",
        "\n",
        "TensorFlow ops like `tf.cond` and `tf.while_loop` continue to work, but control flow is often easier to write and understand when written in Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yCQTtTPTW3WF"
      },
      "outputs": [],
      "source": [
        "# Simple loop\n",
        "\n",
        "@tf.function\n",
        "def f(x):\n",
        "  while tf.reduce_sum(x) > 1:\n",
        "    tf.print(x)\n",
        "    x = tf.tanh(x)\n",
        "  return x\n",
        "\n",
        "f(tf.random.uniform([5]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxwJ8znPI0Cg"
      },
      "source": [
        "If you're curious you can inspect the code autograph generates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jlQD1ffRXJhl"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(f.python_function))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgKmkrNTZSyz"
      },
      "source": [
        "### Conditionals\n",
        "\n",
        "AutoGraph will convert some `if <condition>` statements into the equivalent `tf.cond` calls. This substitution is made if `<condition>` is a Tensor. Otherwise, the `if` statement is executed as a Python conditional.\n",
        "\n",
        "A Python conditional executes during tracing, so exactly one branch of the conditional will be added to the graph. Without AutoGraph, this traced graph would be unable to take the alternate branch if there is data-dependent control flow.\n",
        "\n",
        "`tf.cond` traces and adds both branches of the conditional to the graph, dynamically selecting a branch at execution time. Tracing can have unintended side effects; see [AutoGraph tracing effects](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#effects-of-the-tracing-process) for more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BOQl8PMq2Sf3"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  for i in tf.range(1, n + 1):\n",
        "    print('Tracing for loop')\n",
        "    if i % 15 == 0:\n",
        "      print('Tracing fizzbuzz branch')\n",
        "      tf.print('fizzbuzz')\n",
        "    elif i % 3 == 0:\n",
        "      print('Tracing fizz branch')\n",
        "      tf.print('fizz')\n",
        "    elif i % 5 == 0:\n",
        "      print('Tracing buzz branch')\n",
        "      tf.print('buzz')\n",
        "    else:\n",
        "      print('Tracing default branch')\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(5))\n",
        "fizzbuzz(tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4rBO5AQ15HVC"
      },
      "source": [
        "See the [reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#if-statements) for additional restrictions on AutoGraph-converted if statements."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yho4J0a0ZkQS"
      },
      "source": [
        "### Loops\n",
        "\n",
        "AutoGraph will convert some `for` and `while` statements into the equivalent TensorFlow looping ops, like `tf.while_loop`. If not converted, the `for` or `while` loop is executed as a Python loop.\n",
        "\n",
        "This substitution is made in the following situations:\n",
        "\n",
        "- `for x in y`: if `y` is a Tensor, convert to `tf.while_loop`. In the special case where `y` is a `tf.data.Dataset`, a combination of `tf.data.Dataset` ops are generated.\n",
        "- `while <condition>`: if `<condition>` is a Tensor, convert to `tf.while_loop`.\n",
        "\n",
        "A Python loop executes during tracing, adding additional ops to the `tf.Graph` for every iteration of the loop.\n",
        "\n",
        "A TensorFlow loop traces the body of the loop, and dynamically selects how many iterations to run at execution time.  The loop body only appears once in the generated `tf.Graph`.\n",
        "\n",
        "See the [reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#while-statements) for additional restrictions on AutoGraph-converted `for` and `while` statements."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sp4rbIdfbM6s"
      },
      "source": [
        "#### Looping over Python data\n",
        "\n",
        "A common pitfall is to loop over Python/Numpy data within a `tf.function`. This loop will execute during the tracing process, adding a copy of your model to the `tf.Graph` for each iteration of the loop.\n",
        "\n",
        "If you want to wrap the entire training loop in `tf.function`, the safest way to do this is to wrap your data as a `tf.data.Dataset` so that AutoGraph will dynamically unroll the training loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WGZ19LspbZ27"
      },
      "outputs": [],
      "source": [
        "def measure_graph_size(f, *args):\n",
        "  g = f.get_concrete_function(*args).graph\n",
        "  print(\"{}({}) contains {} nodes in its graph\".format(\n",
        "      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))\n",
        "\n",
        "@tf.function\n",
        "def train(dataset):\n",
        "  loss = tf.constant(0)\n",
        "  for x, y in dataset:\n",
        "    loss += tf.abs(y - x) # Some dummy computation.\n",
        "  return loss\n",
        "\n",
        "small_data = [(1, 1)] * 3\n",
        "big_data = [(1, 1)] * 10\n",
        "measure_graph_size(train, small_data)\n",
        "measure_graph_size(train, big_data)\n",
        "\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: small_data, (tf.int32, tf.int32)))\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: big_data, (tf.int32, tf.int32)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeD2U-yrbfVb"
      },
      "source": [
        "When wrapping Python/Numpy data in a Dataset, be mindful of `tf.data.Dataset.from_generator` versus ` tf.data.Dataset.from_tensors`. The former will keep the data in Python and fetch it via `tf.py_function` which can have performance implications, whereas the latter will bundle a copy of the data as one large `tf.constant()` node in the graph, which can have memory implications.\n",
        "\n",
        "Reading data from files via TFRecordDataset/CsvDataset/etc. is the most effective way to consume data, as then TensorFlow itself can manage the asynchronous loading and prefetching of data, without having to involve Python. To learn more, see the [tf.data guide](../../guide/data)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyksHW9TCukR"
      },
      "source": [
        "#### Accumulating values in a loop\n",
        "\n",
        "A common pattern is to accumulate intermediate values from a loop. Normally, this is accomplished by appending to a Python list or adding entries to a Python dictionary. However, as these are Python side effects, they will not work as expected in a dynamically unrolled loop. Use `tf.TensorArray` to accumulate results from a dynamically unrolled loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HJ3Vb3dXfefN"
      },
      "outputs": [],
      "source": [
        "batch_size = 2\n",
        "seq_len = 3\n",
        "feature_size = 4\n",
        "\n",
        "def rnn_step(inp, state):\n",
        "  return inp + state\n",
        "\n",
        "@tf.function\n",
        "def dynamic_rnn(rnn_step, input_data, initial_state):\n",
        "  # [batch, time, features] -> [time, batch, features]\n",
        "  input_data = tf.transpose(input_data, [1, 0, 2])\n",
        "  max_seq_len = input_data.shape[0]\n",
        "\n",
        "  states = tf.TensorArray(tf.float32, size=max_seq_len)\n",
        "  state = initial_state\n",
        "  for i in tf.range(max_seq_len):\n",
        "    state = rnn_step(input_data[i], state)\n",
        "    states = states.write(i, state)\n",
        "  return tf.transpose(states.stack(), [1, 0, 2])\n",
        "  \n",
        "dynamic_rnn(rnn_step,\n",
        "            tf.random.uniform([batch_size, seq_len, feature_size]),\n",
        "            tf.zeros([batch_size, feature_size]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i2MVoIVaNApG"
      },
      "source": [
        "## Limitations\n",
        "\n",
        "TensorFlow `Function` has a few limitations by design that you should be aware of when converting a Python function to a `Function`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EJqHGFSVLIKl"
      },
      "source": [
        "### Executing Python side effects\n",
        "\n",
        "Side effects, like printing, appending to lists, and mutating globals, can behave unexpectedly inside a `Function`, sometimes executing twice or not all. They only happen the first time you call a `Function` with a set of inputs.  Afterwards, the traced `tf.Graph` is reexecuted, without executing the Python code.\n",
        "\n",
        "The general rule of thumb is to avoid relying on Python side effects in your logic and only use them to debug your traces. Otherwise, TensorFlow APIs like `tf.data`, `tf.print`, `tf.summary`, `tf.Variable.assign`, and `tf.TensorArray` are the best way to ensure your code will be executed by the TensorFlow runtime with each call."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w2sACuZ9TTRk"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  print(\"Traced with\", x)\n",
        "  tf.print(\"Executed with\", x)\n",
        "\n",
        "f(1)\n",
        "f(1)\n",
        "f(2)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1I0dPiqTV8H"
      },
      "source": [
        "If you would like to execute Python code during each invocation of a `Function`, `tf.py_function` is an exit hatch. The drawback of `tf.py_function` is that it's not portable or particularly performant, cannot be saved with SavedModel, and does not work well in distributed (multi-GPU, TPU) setups. Also, since `tf.py_function` has to be wired into the graph, it casts all inputs/outputs to tensors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bOW1v9WVKGgH"
      },
      "source": [
        "#### Changing Python global and free variables\n",
        "\n",
        "Changing Python global and [free variables](https://docs.python.org/3/reference/executionmodel.html#binding-of-names) counts as a Python side effect, so it only happens during tracing.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7aJD--9qTWmg"
      },
      "outputs": [],
      "source": [
        "external_list = []\n",
        "\n",
        "@tf.function\n",
        "def side_effect(x):\n",
        "  print('Python side effect')\n",
        "  external_list.append(x)\n",
        "\n",
        "side_effect(1)\n",
        "side_effect(1)\n",
        "side_effect(1)\n",
        "# The list append only happened once!\n",
        "assert len(external_list) == 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbFG5CX4LwQA"
      },
      "source": [
        "You should avoid mutating containers like lists, dicts, other objects that live outside the `Function`. Instead, use arguments and TF objects. For example, the section [\"Accumulating values in a loop\"](#accumulating_values_in_a_loop) has one example of how list-like operations can be implemented.\n",
        "\n",
        "You can, in some cases, capture and manipulate state if it is a [`tf.Variable`](https://www.tensorflow.org/guide/variable). This is how the weights of Keras models are updated with repeated calls to the same `ConcreteFunction`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X_oNNGrAqPJ1"
      },
      "source": [
        "#### Using Python iterators and generators"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msTmv-oyUNaf"
      },
      "source": [
        "Many Python features, such as generators and iterators, rely on the Python runtime to keep track of state. In general, while these constructs work as expected in eager mode, they are examples of Python side effects and therefore only happen during tracing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FNPD4unZUedH"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def buggy_consume_next(iterator):\n",
        "  tf.print(\"Value:\", next(iterator))\n",
        "\n",
        "iterator = iter([1, 2, 3])\n",
        "buggy_consume_next(iterator)\n",
        "# This reuses the first value from the iterator, rather than consuming the next value.\n",
        "buggy_consume_next(iterator)\n",
        "buggy_consume_next(iterator)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wcS3TAgCjTWR"
      },
      "source": [
        "Just like how TensorFlow has a specialized `tf.TensorArray` for list constructs, it has a specialized `tf.data.Iterator` for iteration constructs. See the section on [AutoGraph Transformations](#autograph_transformations) for an overview. Also, the [`tf.data`](https://www.tensorflow.org/guide/data) API can help implement generator patterns:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8D_iKetXW6VE"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def good_consume_next(iterator):\n",
        "  # This is ok, iterator is a tf.data.Iterator\n",
        "  tf.print(\"Value:\", next(iterator))\n",
        "\n",
        "ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])\n",
        "iterator = iter(ds)\n",
        "good_consume_next(iterator)\n",
        "good_consume_next(iterator)\n",
        "good_consume_next(iterator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHQ0UeU-vWo8"
      },
      "source": [
        "### Deleting tf.Variables between `Function` calls\n",
        "\n",
        "Another error you may encounter is a garbage-collected variable. `ConcreteFunction`s only retain [WeakRefs](https://docs.python.org/3/library/weakref.html) to the variables they close over, so you must retain a reference to any variables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uMiRPfETjpt-"
      },
      "outputs": [],
      "source": [
        "external_var = tf.Variable(3)\n",
        "@tf.function\n",
        "def f(x):\n",
        "  return x * external_var\n",
        "\n",
        "traced_f = f.get_concrete_function(4)\n",
        "print(\"Calling concrete function...\")\n",
        "print(traced_f(4))\n",
        "\n",
        "# The original variable object gets garbage collected, since there are no more\n",
        "# references to it.\n",
        "external_var = tf.Variable(4)\n",
        "print()\n",
        "print(\"Calling concrete function after garbage collecting its closed Variable...\")\n",
        "with assert_raises(tf.errors.FailedPreconditionError):\n",
        "  traced_f(4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-D6nh3QirXAd"
      },
      "source": [
        "## Known Issues\n",
        "\n",
        "If your `Function` is not evaluating correctly, the error may be explained by these known issues which are planned to be fixed in the future."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZoPg5w1Pjqna"
      },
      "source": [
        "### Depending on Python global and free variables\n",
        "\n",
        "`Function` creates a new `ConcreteFunction` when called with a new value of a Python argument. However, it does not do that for the Python closure, globals, or nonlocals of that `Function`. If their value changes in between calls to the `Function`, the `Function` will still use the values they had when it was traced. This is different from how regular Python functions work.\n",
        "\n",
        "For that reason, we recommend a functional programming style that uses arguments instead of closing over outer names."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oeJMdXd3M0cM"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def buggy_add():\n",
        "  return 1 + foo\n",
        "\n",
        "@tf.function\n",
        "def recommended_add(foo):\n",
        "  return 1 + foo\n",
        "\n",
        "foo = 1\n",
        "print(\"Buggy:\", buggy_add())\n",
        "print(\"Correct:\", recommended_add(foo))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L3q7sUJWZOSU"
      },
      "outputs": [],
      "source": [
        "print(\"Updating the value of `foo` to 100!\")\n",
        "foo = 100\n",
        "print(\"Buggy:\", buggy_add())  # Did not change!\n",
        "print(\"Correct:\", recommended_add(foo))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tu0SnPwaL7pI"
      },
      "source": [
        "You can close over outer names, as long as you don't update their values.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvwe9gTIWfx6"
      },
      "source": [
        "#### Depending on Python objects"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJkZS-SwPvOQ"
      },
      "source": [
        "The recommendation to pass Python objects as arguments into `tf.function` has a number of known issues, that are expected to be fixed in the future. In general, you can rely on consistent tracing if you use a Python primitive or `tf.nest`-compatible structure as an argument or pass in a *different* instance of an object into a `Function`. However, `Function` will *not* create a new trace when you pass **the same object and only change its attributes**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ux8KJESVWDxX"
      },
      "outputs": [],
      "source": [
        "class SimpleModel(tf.Module):\n",
        "  def __init__(self):\n",
        "    # These values are *not* tf.Variables.\n",
        "    self.bias = 0.\n",
        "    self.weight = 2.\n",
        "\n",
        "@tf.function\n",
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "simple_model = SimpleModel()\n",
        "x = tf.constant(10.)\n",
        "print(evaluate(simple_model, x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mUxRF4ghZZvX"
      },
      "outputs": [],
      "source": [
        "print(\"Adding bias!\")\n",
        "simple_model.bias += 5.0\n",
        "print(evaluate(simple_model, x))  # Didn't change :("
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ytcgg2qFWaBF"
      },
      "source": [
        "Using the same `Function` to evaluate the updated instance of the model will be buggy since the updated model has the [same cache key](#rules_of_tracing) as the original model.\n",
        "\n",
        "For that reason, we recommend that you write your `Function` to avoid depending on mutable object attributes or create new objects.\n",
        "\n",
        "If that is not possible, one workaround is to make new `Function`s each time you modify your object to force retracing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pFvWmWAAQjrv"
      },
      "outputs": [],
      "source": [
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "new_model = SimpleModel()\n",
        "evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)\n",
        "# Don't pass in `new_model`, `Function` already captured its state during tracing.\n",
        "print(evaluate_no_bias(x))  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bdU2-jF4ZH0B"
      },
      "outputs": [],
      "source": [
        "\n",
        "print(\"Adding bias!\")\n",
        "new_model.bias += 5.0\n",
        "# Create new Function and ConcreteFunction since you modified new_model.\n",
        "evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x)\n",
        "print(evaluate_with_bias(x)) # Don't pass in `new_model`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uFgEZClsZrEi"
      },
      "source": [
        "As [retracing can be expensive](https://www.tensorflow.org/guide/intro_to_graphs#tracing_and_performance), you can use `tf.Variable`s as object attributes, which can be mutated (but not changed, careful!) for a similar effect without needing a retrace.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "daAP_lucwS6w"
      },
      "outputs": [],
      "source": [
        "class BetterModel:\n",
        "\n",
        "  def __init__(self):\n",
        "    self.bias = tf.Variable(0.)\n",
        "    self.weight = tf.Variable(2.)\n",
        "\n",
        "@tf.function\n",
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "better_model = BetterModel()\n",
        "print(evaluate(better_model, x))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ktqwMJBqwTFj"
      },
      "outputs": [],
      "source": [
        "print(\"Adding bias!\")\n",
        "better_model.bias.assign_add(5.0)  # Note: instead of better_model.bias += 5\n",
        "print(evaluate(better_model, x))  # This works!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lPr_6mK_AQWL"
      },
      "source": [
        "### Creating tf.Variables\n",
        "\n",
        "`Function` only supports creating variables once, when first called, and then reusing them. You cannot create `tf.Variables` in new traces. Creating new variables in subsequent calls is currently not allowed, but will be in the future.\n",
        "\n",
        "\n",
        "Example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tx0Vvnb_9OB-"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  v = tf.Variable(1.0)\n",
        "  return v\n",
        "\n",
        "with assert_raises(ValueError):\n",
        "  f(1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KYm6-5GCILXQ"
      },
      "source": [
        "You can create variables inside a `Function` as long as those variables are only created the first time the function is executed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HQrG5_kOiKl_"
      },
      "outputs": [],
      "source": [
        "class Count(tf.Module):\n",
        "  def __init__(self):\n",
        "    self.count = None\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self):\n",
        "    if self.count is None:\n",
        "      self.count = tf.Variable(0)\n",
        "    return self.count.assign_add(1)\n",
        "\n",
        "c = Count()\n",
        "print(c())\n",
        "print(c())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7uD6qI7aJwbR"
      },
      "source": [
        "#### Using with multiple Keras optimizers\n",
        "You may encounter `ValueError: tf.function-decorated function tried to create variables on non-first call.` when using more than one Keras optimizer with a `tf.function`. This error occurs because optimizers internally create `tf.Variables` when they apply gradients for the first time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yWQ3-r99Jvze"
      },
      "outputs": [],
      "source": [
        "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
        "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
        " \n",
        "@tf.function\n",
        "def train_step(w, x, y, optimizer):\n",
        "   with tf.GradientTape() as tape:\n",
        "       L = tf.reduce_sum(tf.square(w*x - y))\n",
        "   gradients = tape.gradient(L, [w])\n",
        "   optimizer.apply_gradients(zip(gradients, [w]))\n",
        "\n",
        "w = tf.Variable(2.)\n",
        "x = tf.constant([-1.])\n",
        "y = tf.constant([2.])\n",
        "\n",
        "train_step(w, x, y, opt1)\n",
        "print(\"Calling `train_step` with different optimizer...\")\n",
        "with assert_raises(ValueError):\n",
        "  train_step(w, x, y, opt2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Q8BRPCThTjB"
      },
      "source": [
        "If you need to change the optimizer during training, a workaround is to create a new `Function` for each optimizer, calling the [`ConcreteFunction`](#obtaining_concrete_functions) directly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YV5F2Gy9hSI3"
      },
      "outputs": [],
      "source": [
        "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
        "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
        "\n",
        "# Not a tf.function.\n",
        "def train_step(w, x, y, optimizer):\n",
        "   with tf.GradientTape() as tape:\n",
        "       L = tf.reduce_sum(tf.square(w*x - y))\n",
        "   gradients = tape.gradient(L, [w])\n",
        "   optimizer.apply_gradients(zip(gradients, [w]))\n",
        "\n",
        "w = tf.Variable(2.)\n",
        "x = tf.constant([-1.])\n",
        "y = tf.constant([2.])\n",
        "\n",
        "# Make a new Function and ConcreteFunction for each optimizer.\n",
        "train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)\n",
        "train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)\n",
        "for i in range(10):\n",
        "  if i % 2 == 0:\n",
        "    train_step_1(w, x, y) # `opt1` is not used as a parameter. \n",
        "  else:\n",
        "    train_step_2(w, x, y) # `opt2` is not used as a parameter."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xjnz5CcuqQac"
      },
      "source": [
        "#### Using with multiple Keras models\n",
        "\n",
        "You may also encounter `ValueError: tf.function-decorated function tried to create variables on non-first call.` when passing different model instances to the same `Function`.\n",
        "\n",
        "This error occurs because Keras models (which [do not have their input shape defined](https://www.tensorflow.org/guide/keras/custom_layers_and_models#best_practice_deferring_weight_creation_until_the_shape_of_the_inputs_is_known)) and Keras layers create `tf.Variables`s when they are first called. You may be attempting to initialize those variables inside a `Function`, which has already been called. To avoid this error, try calling `model.build(input_shape)` to initialize all the weights before training the model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IKyrEY5GVX3M"
      },
      "source": [
        "## Further reading\n",
        "\n",
        "To learn about how to export and load a `Function`, see the [SavedModel guide](../../guide/saved_model). To learn more about graph optimizations that are performed after tracing, see the [Grappler guide](../../guide/graph_optimization). To learn how to optimize your data pipeline and profile your model, see the [Profiler guide](../../guide/profiler.md)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "function.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
